package com.xlh.dokka.svc.core.scheduler.filter;

import com.xlh.dokka.svc.core.model.ContainerConfig;
import com.xlh.dokka.svc.core.model.GpuNode;
import com.xlh.dokka.svc.core.model.Node;
import org.apache.commons.collections.CollectionUtils;
import org.springframework.beans.factory.annotation.Value;

import java.util.ArrayList;
import java.util.List;

/**
 * @author cheer
 */
public class HealthFilter implements Filter {

    public HealthFilter(Boolean recorded) {
        this.recorded = recorded;
    }

    private Boolean recorded;

    @Override
    public List<Node> filter(ContainerConfig config, List<Node> nodes) {
        List<Node> result = new ArrayList<>();
        if (CollectionUtils.isNotEmpty(nodes)) {
            for (Node node : nodes) {
                if (node.getHealthy() != null && node.getHealthy() && node.getNodeUsage() != null &&
                        node.getResourceUsage() != null) {
                    if (recorded) {
                        if (!node.getGpu()) {
                            result.add(node);
                        }
                    } else {
                        result.add(node);
                    }
                }
            }
        }
        return result;
    }

    @Override
    public List<GpuNode> gpuFilter(ContainerConfig config, List<GpuNode> gpuNodes) {
        List<GpuNode> result = new ArrayList<>();
        if (CollectionUtils.isNotEmpty(gpuNodes)) {
            for (GpuNode gpuNode : gpuNodes) {
                Node node = gpuNode.getNode();
                if (node.getHealthy() != null && node.getHealthy() && node.getNodeUsage() != null &&
                        node.getResourceUsage() != null) {
                    result.add(gpuNode);
                }
            }
        }
        return result;
    }
}
